from pymongo import UpdateOne
from zc_core.pipelines.helper.catalog_helper import CatalogHelper
from zc_core.client.mongo_client import Mongo

mongo = Mongo()


def fix_data_catalog():
    update_bulk = list()
    src_list = mongo.list('data_20210702',
                          fields={'_id': 1, 'catalog1Id': 1, 'catalog1Name': 1, 'catalog1Code': 1, 'catalog2Id': 1,
                                  'catalog2Name': 1, 'catalog2Code': 1,
                                  'catalog3Id': 1, 'catalog3Name': 1, 'catalog3Code': 1})
    for item in src_list:
        _id = item.get('_id')
        catalog1Id = item.get('catalog1Id')
        catalog1Name = item.get('catalog1Name')
        catalog1Code = item.get('catalog1Code')
        catalog2Id = item.get('catalog2Id')
        catalog2Name = item.get('catalog2Name')
        catalog2Code = item.get('catalog2Code')
        catalog3Id = item.get('catalog3Id')
        catalog3Name = item.get('catalog3Name')
        catalog3Code = item.get('catalog3Code')
        if catalog1Id:
            update_bulk.append(UpdateOne({'_id': _id}, {'$set': {
                '_id': _id,
                'catalog1Id': catalog1Id,
                'catalog1Name': catalog1Name,
                'catalog1Code': catalog1Code,
                'catalog2Id': catalog2Id,
                'catalog2Name': catalog2Name,
                'catalog2Code': catalog2Code,
                'catalog3Id': catalog3Id,
                'catalog3Name': catalog3Name,
                'catalog3Code': catalog3Code,
            }}, upsert=False))

    print(len(update_bulk))
    mongo.bulk_write('data_20210709', update_bulk)
    mongo.close()
    print('任务完成~')


# 从sku_pool池子修补全部分类
def fix_sku_pool_catalog():
    update_bulk = list()
    src_list = mongo.list('sku_pool', query={"catalog1Id": {"$exists": False}}, fields={'_id': 1, 'catalog3Id': 1})
    for item in src_list:
        _id = item.get('_id')
        catalog3Id = item.get('catalog3Id')
        dict1 = {}
        dict1['_id'] = _id
        dict1['catalog3Id'] = catalog3Id
        if catalog3Id:
            try:
                cat_helper = CatalogHelper()
                cat_helper.fill(dict1)
                update_bulk.append(UpdateOne({'_id': _id}, {'$set': dict1}, upsert=False))
            except Exception as e:
                print(e)
    print(len(update_bulk))
    mongo.bulk_write('sku_pool', update_bulk)
    list1 = mongo.list('sku_pool', query={"catalog1Id": {"$exists": False}}, fields={'_id': 1, 'catalog3Id': 1})
    print('剩余未有分类不完整的数量', len(list1))
    mongo.close()
    print('任务完成~')


if __name__ == '__main__':
    # fix_data_catalog()
    fix_sku_pool_catalog()
